package com.example.myapplication;

import android.content.Context;
import android.graphics.Bitmap;
import android.graphics.BitmapFactory;
import android.graphics.ImageFormat;
import android.graphics.Matrix;
import android.graphics.Rect;
import android.graphics.YuvImage;
import android.media.Image;
import android.media.Image.Plane;
import android.media.ImageReader;
import android.os.SystemClock;
import android.util.Log;
import androidx.annotation.NonNull;
import androidx.annotation.VisibleForTesting;
import androidx.camera.core.ImageProxy;
import com.google.mediapipe.framework.image.BitmapImageBuilder;
import com.google.mediapipe.framework.image.MPImage;
import com.google.mediapipe.tasks.core.BaseOptions;
import com.google.mediapipe.tasks.core.Delegate;
import com.google.mediapipe.tasks.vision.core.RunningMode;
import com.google.mediapipe.tasks.vision.handlandmarker.HandLandmarker;
import com.google.mediapipe.tasks.vision.handlandmarker.HandLandmarkerResult;
import com.google.mediapipe.tasks.vision.handlandmarker.HandLandmarker.HandLandmarkerOptions;

import java.io.ByteArrayOutputStream;
import java.nio.ByteBuffer;
import java.util.Optional;

public class HandLandmarkerHelper {
    private static final String TAG = "HandLandmarkerHelper";
    private static final String MP_HAND_LANDMARKER_TASK = "hand_landmarker.task";

    public static final int DELEGATE_CPU = 0;
    public static final int DELEGATE_GPU = 1;
    public static final float DEFAULT_HAND_DETECTION_CONFIDENCE = 0.5F;
    public static final float DEFAULT_HAND_TRACKING_CONFIDENCE = 0.5F;
    public static final float DEFAULT_HAND_PRESENCE_CONFIDENCE = 0.5F;
    public static final int DEFAULT_NUM_HANDS = 2;
    public static final int OTHER_ERROR = 0;
    public static final int GPU_ERROR = 1;
    private float minHandDetectionConfidence = DEFAULT_HAND_DETECTION_CONFIDENCE;
    private float minHandTrackingConfidence = DEFAULT_HAND_TRACKING_CONFIDENCE;
    private float minHandPresenceConfidence = DEFAULT_HAND_PRESENCE_CONFIDENCE;
    private int maxNumHands = DEFAULT_NUM_HANDS;
    private int currentDelegate = DELEGATE_GPU;
    private RunningMode runningMode = RunningMode.IMAGE;
    private final Context context;
    private final LandmarkerListener handLandmarkerHelperListener;
    private HandLandmarker handLandmarker;

    public HandLandmarkerHelper(float minHandDetectionConfidence, float minHandTrackingConfidence,
                                float minHandPresenceConfidence, int maxNumHands, int currentDelegate,
                                RunningMode runningMode, Context context, LandmarkerListener handLandmarkerHelperListener) {

        this.minHandDetectionConfidence = minHandDetectionConfidence;
        this.minHandTrackingConfidence = minHandTrackingConfidence;
        this.minHandPresenceConfidence = minHandPresenceConfidence;
        this.maxNumHands = maxNumHands;
        this.currentDelegate = currentDelegate;
        this.runningMode = runningMode;
        this.context = context;
        this.handLandmarkerHelperListener = handLandmarkerHelperListener;

        setupHandLandmarker();
    }

    public void clearHandLandmarker() {
        if (handLandmarker != null) {
            handLandmarker.close();
            handLandmarker = null;
        }
    }

    public boolean isClose() {
        return handLandmarker == null;
    }

    private void setupHandLandmarker() {
        BaseOptions baseOptions = BaseOptions.builder()
                .setDelegate(Delegate.GPU)
                .setModelAssetPath(MP_HAND_LANDMARKER_TASK)
                .build();

        HandLandmarkerOptions.Builder optionsBuilder = HandLandmarkerOptions.builder()
                .setBaseOptions(baseOptions)
                .setMinHandDetectionConfidence(minHandDetectionConfidence)
                .setMinHandPresenceConfidence(minHandPresenceConfidence)
                .setMinTrackingConfidence(minHandTrackingConfidence)
                .setNumHands(maxNumHands)
                .setRunningMode(runningMode);

        if (runningMode == RunningMode.LIVE_STREAM) {
            optionsBuilder.setResultListener(this::returnLivestreamResult)
                    .setErrorListener(this::returnLivestreamError);
        }

        handLandmarker = HandLandmarker.createFromOptions(context, optionsBuilder.build());
    }

    public void detectLiveStream(@NonNull ImageProxy imageProxy, boolean isFrontCamera) {
        if (runningMode != RunningMode.LIVE_STREAM) {
            throw new IllegalArgumentException("Attempting to call detectLiveStream while not using RunningMode.LIVE_STREAM");
        }

        long frameTime = SystemClock.uptimeMillis();

        Bitmap bitmap = yuv420ToBitmap(imageProxy);

        Matrix matrix = new Matrix();
        matrix.postRotate(imageProxy.getImageInfo().getRotationDegrees());
        if (isFrontCamera) {
            matrix.postScale(-1f, 1f, imageProxy.getWidth() / 2f, imageProxy.getHeight() / 2f);
        }
        Bitmap rotatedBitmap = Bitmap.createBitmap(bitmap, 0, 0, bitmap.getWidth(), bitmap.getHeight(), matrix, true);

        MPImage mpImage = new BitmapImageBuilder(rotatedBitmap).build();

        handLandmarker.detectAsync(mpImage, frameTime);

        imageProxy.close();
    }

    private Bitmap yuv420ToBitmap(ImageProxy image) {

        //转换为BITMAP进行推理
        ImageProxy.PlaneProxy[] planes = image.getPlanes();
        ByteBuffer yBuffer = planes[0].getBuffer();
        ByteBuffer uBuffer = planes[1].getBuffer();
        ByteBuffer vBuffer = planes[2].getBuffer();

        int ySize = yBuffer.remaining();
        int uSize = uBuffer.remaining();
        int vSize = vBuffer.remaining();

        byte[] nv21 = new byte[ySize + uSize + vSize];

        yBuffer.get(nv21, 0, ySize);
        vBuffer.get(nv21, ySize, vSize);
        uBuffer.get(nv21, ySize + vSize, uSize);

        YuvImage yuvImage = new YuvImage(nv21, ImageFormat.NV21, image.getWidth(), image.getHeight(), null);
        ByteArrayOutputStream out = new ByteArrayOutputStream();
        yuvImage.compressToJpeg(new Rect(0, 0, image.getWidth(), image.getHeight()), 100, out);
        byte[] imageBytes = out.toByteArray();

        return BitmapFactory.decodeByteArray(imageBytes, 0, imageBytes.length);
    }

    @VisibleForTesting
    public void detectAsync(MPImage mpImage, long frameTime) {
        handLandmarker.detectAsync(mpImage, frameTime);
    }

    private void returnLivestreamResult(HandLandmarkerResult result, MPImage input) {
        long finishTimeMs = SystemClock.uptimeMillis();
        long inferenceTime = finishTimeMs - result.timestampMs();

        handLandmarkerHelperListener.onResults(new ResultBundle(Optional.of(result), inferenceTime, input.getHeight(), input.getWidth()));
    }

    private void returnLivestreamError(RuntimeException error) {
        handLandmarkerHelperListener.onError(error.getMessage() != null ? error.getMessage() : "An unknown error has occurred", GPU_ERROR);
    }

    public interface LandmarkerListener {
        void onError(String error, int errorCode);

        void onResults(ResultBundle resultBundle);
    }

    public static class ResultBundle {
        private final Optional<HandLandmarkerResult> result;
        private final long inferenceTime;
        private final int inputImageHeight;
        private final int inputImageWidth;

        public ResultBundle(Optional<HandLandmarkerResult> result, long inferenceTime, int inputImageHeight, int inputImageWidth) {
            this.result = result;
            this.inferenceTime = inferenceTime;
            this.inputImageHeight = inputImageHeight;
            this.inputImageWidth = inputImageWidth;
        }

        public Optional<HandLandmarkerResult> getResult() {
            return result;
        }

        public long getInferenceTime() {
            return inferenceTime;
        }

        public int getInputImageHeight() {
            return inputImageHeight;
        }

        public int getInputImageWidth() {
            return inputImageWidth;
        }
    }
}
